package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

// Config 定义配置结构
type Config struct {
	WebhookURL           string  `json:"webhook_url"`
	ExpectedGPUCount     int     `json:"expected_gpu_count"`
	CheckInterval        int     `json:"check_interval"`
	TemperatureThreshold int     `json:"temperature_threshold"`
	MemoryUsageThreshold float64 `json:"memory_usage_threshold"`
	PowerThreshold       int     `json:"power_threshold"`
}

// AlertMessage 定义告警消息结构
type AlertMessage struct {
	MsgType string `json:"msg_type"`
	Content struct {
		Text string `json:"text"`
	} `json:"content"`
}

// GPUStatus 定义GPU状态结构
type GPUStatus struct {
	Count    int
	Errors   []string
	HasError bool
}

// 全局变量存储上一次的GPU状态和配置
var (
	lastGPUStatus GPUStatus
	config        Config
)

// 全局变量存储主机信息
var (
	hostname   string
	externalIP string
)

// loadConfig 加载配置文件
func loadConfig() error {
	// 获取可执行文件的路径
	exePath, err := os.Executable()
	if err != nil {
		return fmt.Errorf("获取可执行文件路径失败: %v", err)
	}

	// 获取可执行文件所在目录
	exeDir := filepath.Dir(exePath)

	// 构造配置文件的完整路径
	configPath := filepath.Join(exeDir, "config.json")

	// 读取配置文件
	data, err := ioutil.ReadFile(configPath)
	if err != nil {
		return fmt.Errorf("读取配置文件失败: %v", err)
	}

	// 解析JSON
	if err := json.Unmarshal(data, &config); err != nil {
		return fmt.Errorf("解析配置文件失败: %v", err)
	}

	// 验证配置
	if config.WebhookURL == "" {
		return fmt.Errorf("Webhook URL不能为空")
	}

	if config.CheckInterval <= 0 {
		config.CheckInterval = 30 // 默认30秒
	}

	if config.TemperatureThreshold <= 0 {
		config.TemperatureThreshold = 85 // 默认85度
	}

	if config.MemoryUsageThreshold <= 0 {
		config.MemoryUsageThreshold = 95 // 默认95%
	}

	if config.PowerThreshold <= 0 {
		config.PowerThreshold = 200000 // 默认200W (单位: mW)
	}

	return nil
}

// getGPUCount 获取NVIDIA GPU数量
func getGPUCount() (int, error) {
	// 使用nvidia-smi命令获取GPU数量
	cmd := exec.Command("nvidia-smi", "--list-gpus")
	var stdout bytes.Buffer
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		return 0, fmt.Errorf("执行nvidia-smi命令失败: %v", err)
	}

	count := 0
	for {
		_, err := stdout.ReadString('\n')
		if err != nil {
			break
		}
		count++
	}

	return count, nil
}

// checkGPUDetails 检查GPU详细状态
func checkGPUDetails() ([]string, error) {
	// 使用nvidia-smi检查GPU详细状态
	cmd := exec.Command("nvidia-smi")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	if err != nil {
		return nil, fmt.Errorf("执行nvidia-smi命令失败: %v, stderr: %s", err, stderr.String())
	}

	// 解析输出检查是否有异常
	output := stdout.String()
	var errors []string

	// 检查输出中是否包含错误信息
	if strings.Contains(output, "ERR!") {
		// 提取包含ERR!的行
		lines := strings.Split(output, "\n")
		for _, line := range lines {
			if strings.Contains(line, "ERR!") {
				errors = append(errors, fmt.Sprintf("GPU状态异常: %s", strings.TrimSpace(line)))
			}
		}
	}

	// 检查其他可能的错误
	if strings.Contains(output, "Error") || strings.Contains(output, "error") {
		errors = append(errors, "GPU监控发现错误")
	}

	return errors, nil
}

// checkGPUProcesses 检查GPU进程状态
func checkGPUProcesses() ([]string, error) {
	// 使用nvidia-smi检查GPU进程
	cmd := exec.Command("nvidia-smi", "pmon", "-c", "1")
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	err := cmd.Run()

	if err != nil {
		// 如果命令执行失败，可能是没有进程在运行，这不是错误
		if stderr.Len() > 0 {
			return nil, fmt.Errorf("执行nvidia-smi pmon命令失败: %v, stderr: %s", err, stderr.String())
		}
	}

	// 解析输出检查是否有异常
	output := stdout.String()
	var errors []string

	// 检查输出中是否包含错误信息
	if strings.Contains(output, "ERR!") || strings.Contains(output, "Error") {
		errors = append(errors, fmt.Sprintf("GPU进程监控发现错误: %s", output))
	}

	return errors, nil
}

// sendFeishuAlert 发送飞书告警
func sendFeishuAlert(message string) error {
	// 构造飞书消息
	alert := AlertMessage{
		MsgType: "text",
	}
	alert.Content.Text = message

	// 转换为JSON
	jsonData, err := json.Marshal(alert)
	if err != nil {
		return fmt.Errorf("序列化消息失败: %v", err)
	}

	// 发送POST请求到飞书Webhook
	resp, err := http.Post(config.WebhookURL, "application/json", strings.NewReader(string(jsonData)))
	if err != nil {
		return fmt.Errorf("发送请求失败: %v", err)
	}
	defer resp.Body.Close()

	// 读取响应
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return fmt.Errorf("读取响应失败: %v", err)
	}

	// 检查响应状态
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("请求失败，状态码: %d, 响应: %s", resp.StatusCode, string(body))
	}

	log.Printf("飞书告警发送成功: %s", string(body))
	return nil
}

// getHostname 获取主机名
func getHostname() (string, error) {
	return os.Hostname()
}

// getExternalIP 获取外网IP地址
func getExternalIP() (string, error) {
	// 尝试从多个公共服务获取外网IP
	services := []string{
		"https://api.ipify.org",
		"https://icanhazip.com",
		"https://ident.me",
		"https://ipecho.net/plain",
	}

	client := &http.Client{
		Timeout: 10 * time.Second,
	}

	for _, service := range services {
		resp, err := client.Get(service)
		if err != nil {
			continue
		}
		defer resp.Body.Close()

		if resp.StatusCode == http.StatusOK {
			body, err := ioutil.ReadAll(resp.Body)
			if err != nil {
				continue
			}
			ip := strings.TrimSpace(string(body))
			if net.ParseIP(ip) != nil {
				return ip, nil
			}
		}
	}

	return "", fmt.Errorf("无法获取外网IP地址")
}

// initSystemInfo 初始化系统信息
func initSystemInfo() error {
	var err error

	// 获取主机名
	hostname, err = getHostname()
	if err != nil {
		hostname = "unknown"
		log.Printf("获取主机名失败: %v", err)
	}

	// 获取外网IP
	externalIP, err = getExternalIP()
	if err != nil {
		externalIP = "unknown"
		log.Printf("获取外网IP失败: %v", err)
	}

	log.Printf("系统信息 - 主机名: %s, 外网IP: %s", hostname, externalIP)
	return nil
}

// monitorGPUs 监控GPU状态
func monitorGPUs() {
	log.Println("开始监控NVIDIA GPU状态...")

	for {
		currentStatus := GPUStatus{}
		hasError := false

		// 获取GPU数量
		count, err := getGPUCount()
		if err != nil {
			currentStatus.Errors = append(currentStatus.Errors, fmt.Sprintf("获取GPU数量失败: %v", err))
			hasError = true
		} else {
			currentStatus.Count = count

			// 检查GPU数量是否符合预期
			if count != config.ExpectedGPUCount {
				currentStatus.Errors = append(currentStatus.Errors, fmt.Sprintf("GPU数量异常: 期望 %d 个, 实际 %d 个", config.ExpectedGPUCount, count))
				hasError = true
			}
		}

		// 检查GPU详细状态
		detailErrors, err := checkGPUDetails()
		if err != nil {
			currentStatus.Errors = append(currentStatus.Errors, fmt.Sprintf("检查GPU详细状态失败: %v", err))
			hasError = true
		} else if len(detailErrors) > 0 {
			currentStatus.Errors = append(currentStatus.Errors, detailErrors...)
			hasError = true
		}

		// 检查GPU进程状态
		processErrors, err := checkGPUProcesses()
		if err != nil {
			currentStatus.Errors = append(currentStatus.Errors, fmt.Sprintf("检查GPU进程失败: %v", err))
			hasError = true
		} else if len(processErrors) > 0 {
			currentStatus.Errors = append(currentStatus.Errors, processErrors...)
			hasError = true
		}

		currentStatus.HasError = hasError

		// 只有在状态发生变化且出现错误时才发送告警
		// 如果从正常状态变为错误状态，发送告警
		// 如果从错误状态变为正常状态，不发送通知（根据用户要求）
		if currentStatus.HasError && !lastGPUStatus.HasError {
			// 从正常状态变为错误状态，发送告警
			message := fmt.Sprintf("🚨 NVIDIA GPU监控告警 🚨\n主机名: %s\n外网IP: %s\n时间: %s\n错误信息:\n",
				hostname, externalIP, time.Now().Format("2006-01-02 15:04:05"))
			for _, err := range currentStatus.Errors {
				message += fmt.Sprintf("- %s\n", err)
			}

			log.Printf("检测到GPU异常，发送告警: %s", message)
			if err := sendFeishuAlert(message); err != nil {
				log.Printf("发送飞书告警失败: %v", err)
			}
		} else if !currentStatus.HasError && lastGPUStatus.HasError {
			// 从错误状态变为正常状态，根据用户要求不发送通知
			log.Printf("GPU状态已恢复正常，但不发送通知")
		}

		// 更新上一次的状态
		lastGPUStatus = currentStatus

		// 等待下一次检查
		time.Sleep(time.Duration(config.CheckInterval) * time.Second)
	}
}

func main() {
	// 加载配置
	if err := loadConfig(); err != nil {
		log.Fatalf("加载配置失败: %v", err)
	}

	// 初始化系统信息
	if err := initSystemInfo(); err != nil {
		log.Printf("初始化系统信息失败: %v", err)
	}

	log.Println("NVIDIA GPU监控程序启动")
	log.Printf("期望GPU数量: %d", config.ExpectedGPUCount)
	log.Printf("检查间隔: %d 秒", config.CheckInterval)

	// 启动监控
	monitorGPUs()
}
